package com.cheems.cheemsspringaialibaba.Tools;

import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.annotation.Tool;
import org.springframework.ai.tool.annotation.ToolParam;

import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;

@Slf4j
public class SearchFromDify {
    //@Tool(description = "信息化业务工作流", returnDirect = true)
    //String SearchInfoWorkFlow(@ToolParam(description = "咨询内容") String message, ToolContext toolContext) throws IOException, InterruptedException {
//
    //    HttpClient client = HttpClient.newHttpClient();
    //    URI uri = URI.create("http://nat.soarli.top:36009/v1/chat-messages");
    //    log.info("检索内容: {}", message);
//
    //    // 构建请求体 - 根据新的API格式
    //    JSONObject requestBody = new JSONObject();
    //    requestBody.put("inputs", new JSONObject()); // 空的inputs对象
    //    requestBody.put("query", message);
    //    requestBody.put("response_mode", "streaming");
    //    requestBody.put("conversation_id", "");
    //    requestBody.put("user", "abc-123");
//
    //    // 添加files数组（如果需要的话，这里暂时留空）
    //    requestBody.put("files", new JSONObject[0]);
//
    //    try {
    //        log.info("请求体JSON: {}", requestBody.toJSONString());
//
    //        HttpRequest request = HttpRequest.newBuilder()
    //                .uri(uri)
    //                .header("Authorization", "Bearer app-KJPssAUaLoGZMo8dzpvgqsiD")
    //                .header("Content-Type", "application/json")
    //                .header("Accept", "*/*")
    //                .POST(HttpRequest.BodyPublishers.ofString(requestBody.toJSONString()))
    //                .build();
//
    //        log.info("HTTP请求已构建");
    //        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
    //        log.info("响应状态码: {}", response.statusCode());
    //        log.info("响应体: {}", response.body());
//
    //        // 解析响应
    //        String responseBody = response.body();
    //        JSONObject jsonResponse = JSONObject.parseObject(responseBody);
//
    //        StringBuilder result = new StringBuilder("查询结果：\n\n");
//
    //        // 添加查询内容
    //        result.append("您的问题: ").append(message).append("\n\n");
//
    //        // 根据新API的响应格式解析结果
    //        // 由于response_mode是streaming，可能需要处理不同的响应格式
    //        if (jsonResponse.containsKey("answer")) {
    //            result.append("找到以下信息：\n")
    //                    .append(jsonResponse.getString("answer"))
    //                    .append("\n\n");
    //        } else if (jsonResponse.containsKey("data")) {
    //            // 如果有data字段，尝试解析
    //            Object data = jsonResponse.get("data");
    //            result.append("找到以下信息：\n")
    //                    .append(data.toString())
    //                    .append("\n\n");
    //        } else {
    //            // 如果格式不明确，返回原始响应
    //            result.append("API响应：\n")
    //                    .append(responseBody)
    //                    .append("\n\n");
    //        }
//
    //        return result.toString();
    //    } catch (Exception e) {
    //        log.error("执行检索工具时发生错误: ", e);
    //        return "检索过程中出现错误: " + e.getMessage();
    //    }
    //}

    @Tool(description = "查询信息化知识库")
    String SearchInfoRag(@ToolParam(description = "咨询内容") String message) throws IOException, InterruptedException {

        HttpClient client = HttpClient.newHttpClient();
        URI uri = URI.create("http://nat.soarli.top:36009/v1/datasets/4a77a121-6529-4626-98e4-812b53156e7a/retrieve");
        log.info("检索内容: {}", message);

        // 构建请求体
        JSONObject requestBody = new JSONObject();
        requestBody.put("query", message);

        // 构建retrieval_model对象
        JSONObject retrievalModel = new JSONObject();
        retrievalModel.put("search_method", "keyword_search");
        retrievalModel.put("reranking_enable", false);
        retrievalModel.put("reranking_mode", "");

        JSONObject rerankingModel = new JSONObject();
        rerankingModel.put("reranking_provider_name", "");
        rerankingModel.put("reranking_model_name", "");
        retrievalModel.put("reranking_model", rerankingModel);

        retrievalModel.put("weights", "");
        retrievalModel.put("top_k", 3);
        retrievalModel.put("score_threshold_enabled", false);
        retrievalModel.put("score_threshold", "");

        requestBody.put("retrieval_model", retrievalModel);

            log.info("请求体JSON: {}", requestBody.toJSONString());

            HttpRequest request = HttpRequest.newBuilder()
                    .uri(uri)
                    .header("Authorization", "Bearer dataset-GisVcn2s7KNcWEhpcKclfA62")
                    .header("Content-Type", "application/json")
                    .POST(HttpRequest.BodyPublishers.ofString(requestBody.toJSONString()))
                    .build();

            log.info("HTTP请求已构建");
            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            log.info("响应状态码: {}", response.statusCode());
            log.info("响应体: {}", response.body());

            // 解析响应
            String responseBody = response.body();
            JSONObject jsonResponse = JSONObject.parseObject(responseBody);

            StringBuilder result = new StringBuilder("查询结果：\n\n");

            // 添加查询内容
            if (jsonResponse.containsKey("query") && jsonResponse.getJSONObject("query").containsKey("content")) {
                result.append("您的问题: ").append(jsonResponse.getJSONObject("query").getString("content")).append("\n\n");
            }

            // 提取记录内容
            if (jsonResponse.containsKey("records") && !jsonResponse.getJSONArray("records").isEmpty()) {
                for (int i = 0; i < jsonResponse.getJSONArray("records").size(); i++) {
                    JSONObject record = jsonResponse.getJSONArray("records").getJSONObject(i);
                    if (record.containsKey("segment") && record.getJSONObject("segment").containsKey("content")) {
                        JSONObject segment = record.getJSONObject("segment");
                        result.append("找到以下信息：\n")
                                .append(segment.getString("content"))
                                .append("\n\n");

                        // 添加文档信息
                        if (segment.containsKey("document") && segment.getJSONObject("document").containsKey("name")) {
                            result.append("来源文档: ").append(segment.getJSONObject("document").getString("name"));
                        }
                    }
                }
            } else {
                result.append("未找到相关信息。");
            }

            return result.toString();

    }

    @Tool(description = "查询财务处知识库")
    String retrieveFromDataset(@ToolParam(description = "咨询内容") String message,ToolContext toolContext) throws IOException, InterruptedException {

        HttpClient client = HttpClient.newHttpClient();
        URI uri = URI.create("https://api.dify.ai/v1/datasets/69f4473f-1ac5-483e-9906-be1f61dcc79b/retrieve");
        log.info("检索内容: {}", message);

        // 构建请求体
        JSONObject requestBody = new JSONObject();
        requestBody.put("query", message);

        // 构建retrieval_model对象
        JSONObject retrievalModel = new JSONObject();
        retrievalModel.put("search_method", "keyword_search");
        retrievalModel.put("reranking_enable", false);
        retrievalModel.put("reranking_mode", "");

        JSONObject rerankingModel = new JSONObject();
        rerankingModel.put("reranking_provider_name", "");
        rerankingModel.put("reranking_model_name", "");
        retrievalModel.put("reranking_model", rerankingModel);

        retrievalModel.put("weights", "");
        retrievalModel.put("top_k", 1);
        retrievalModel.put("score_threshold_enabled", false);
        retrievalModel.put("score_threshold", "");

        requestBody.put("retrieval_model", retrievalModel);

        try {
            log.info("请求体JSON: {}", requestBody.toJSONString());

            HttpRequest request = HttpRequest.newBuilder()
                    .uri(uri)
                    .header("Authorization", "Bearer dataset-80CNBDP8vdorljcHM7oAbSk9")
                    .header("Content-Type", "application/json")
                    .POST(HttpRequest.BodyPublishers.ofString(requestBody.toJSONString()))
                    .build();

            log.info("HTTP请求已构建");
            HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
            log.info("响应状态码: {}", response.statusCode());
            log.info("响应体: {}", response.body());

            // 解析响应
            String responseBody = response.body();
            JSONObject jsonResponse = JSONObject.parseObject(responseBody);

            StringBuilder result = new StringBuilder("查询结果：\n\n");

            // 添加查询内容
            if (jsonResponse.containsKey("query") && jsonResponse.getJSONObject("query").containsKey("content")) {
                result.append("您的问题: ").append(jsonResponse.getJSONObject("query").getString("content")).append("\n\n");
            }

            // 提取记录内容
            if (jsonResponse.containsKey("records") && !jsonResponse.getJSONArray("records").isEmpty()) {
                for (int i = 0; i < jsonResponse.getJSONArray("records").size(); i++) {
                    JSONObject record = jsonResponse.getJSONArray("records").getJSONObject(i);
                    if (record.containsKey("segment") && record.getJSONObject("segment").containsKey("content")) {
                        JSONObject segment = record.getJSONObject("segment");
                        result.append("找到以下信息：\n")
                                .append(segment.getString("content"))
                                .append("\n\n");

                        // 添加文档信息
                        if (segment.containsKey("document") && segment.getJSONObject("document").containsKey("name")) {
                            result.append("来源文档: ").append(segment.getJSONObject("document").getString("name"));
                        }
                    }
                }
            } else {
                result.append("未找到相关信息。");
            }

            return result.toString();
        } catch (Exception e) {
            log.error("执行检索工具时发生错误: ", e);
            return "检索过程中出现错误: " + e.getMessage();
        }
    }
}
